from transformers import BertTokenizer, BertForSequenceClassification
import torch
import json

def process_data(data):
    # 清洗数据，排除第一个元素
    cleaned_data = [d for i, d in enumerate(data) if i != 0]
    # print("Process_data中的历史记录: ")
    # print(cleaned_data)
    # 加载微调后的模型权重和配置文件
    model_path = 'fine_tuned_bert_base_chinese_7'  # 你保存模型的路径
    model = BertForSequenceClassification.from_pretrained(model_path)

    # 加载预训练的 tokenizer
    tokenizer = BertTokenizer.from_pretrained('bert-base-chinese')

    text_groups = []
    user_contents = [d['content'] for d in cleaned_data if d['role'] == 'user']

    json_data = [
        ["我是个体户", "我不是个体户"],
        ["25人以下", "25人以上100人以下", "100以上500人以下", "500以上1000以下", "1000人以上"],
        ["不是一家科技公司", "是一家科技公司"],
        ["提供精神健康福利作为医疗保险", "不提供精神健康福利作为医疗保险", "不符合投保条件", "我不知道是否提供精神健康福利作为医疗保险"],
        ["我知道提供的保险中有哪些心理健康保健的选择", "不，我不知道提供的保险中有哪些心理健康保健的选择", "我不确定提供的保险中有哪些心理健康保健的选择"],
        ["讨论过心理健康问题", "没有讨论过心理健康问题", "我不确定是否讨论过心理健康问题"],
        ["提供了资源", "没有提供资源", "我不确定是否提供资源"],
        ["我的匿名性受到保护", "我的匿名性没有受到保护", "我不确定我的匿名性受到保护"],
        ["请病假会非常简单", "请病假可能会简单", "请病假难度不变", "请病假稍微困难", "请病假非常困难", "我不知道请病假是否会被影响"],
        ["和雇主讨论精神疾病会有负面影响", "和雇主讨论精神疾病不会有负面影响", "我不确定和雇主讨论精神疾病会有负面影响"],
        ["和雇主讨论身体疾病会有负面影响", "和雇主讨论身体疾病不会有负面影响", "我不确定和雇主讨论身体疾病会有负面影响"],
        ["愿意和同事讨论精神疾病", "不愿意和同事讨论精神疾病", "可能愿意和同事讨论精神疾病"],
        ["愿意和直接主管讨论心理健康障碍", "不愿意和直接主管讨论心理健康障碍", "我可能愿意和直接主管讨论心理健康障碍"],
        ["雇主像对待身体健康一样重视心理健康", "雇主不像对待身体健康一样重视心理健康", "不知道雇主是否像对待身体健康一样重视心理健康"],
        ["听说过或者观察到在工作场所公开谈论心理健康问题的同事所带来的负面影响", "不，我没有听说过或观察到在工作场所公开谈论心理健康问题的同事所带来的负面影响"],
        ["我以前有雇主", "我以前没有雇主"],
        ["以前的雇主提供心理健康福利", "以前的雇主没有提供心理健康福利", "部分以前的雇主提供过心理健康福利", "我不知道以前的雇主是否提供心理健康福利"],
        ["我知道以前的雇主提供的心理健康护理方案", "我知道有些以前的雇主提供的心理健康护理方案", "我目前不知道以前的雇主提供的心理健康护理方案", "我到后来才意识到以前的雇主提供的心理健康护理方案"],
        ["以前的雇主有正式讨论过心理健康问题", "以前的雇主没有正式讨论过心理健康问题", "以前的部分雇主正式讨论过心理健康问题", "我不知道以前的雇主是否正式讨论过心理健康问题"],
        ["以前的雇主提供资源", "以前的雇主没有提供资源", "以前的部分雇主提供资源", "我不知道以前的雇主是否提供资源"],
        ["匿名性得到保护", "匿名性没有得到保护", "匿名性有时得到保护", "我不知道匿名性是否得到保护"],
        ["和以前的雇主讨论精神疾病有负面影响", "和以前的雇主讨论精神疾病没有负面影响", "和以前的部分雇主讨论精神疾病有负面影响", "我不知道和以前的雇主讨论精神疾病是否有负面精神"],
        ["和以前的雇主讨论身体疾病有负面影响", "和以前的雇主讨论身体疾病没有负面影响", "和以前的部分雇主讨论身体疾病有负面影响"],
        ["愿意和以前的同事讨论心理健康问题", "不愿意和以前的同事讨论心理健康问题", "愿意和以前的部分同事讨论心理健康问题"],
        ["愿意和直接上司讨论心理健康问题", "不愿意和直接上司讨论心理健康问题", "愿意和以前的部分直接上司讨论心理健康问题", "我不知道是否愿意和直接上司讨论心理健康问题"],
        ["以前的雇主把心理健康看得和身体健康一样重要", "以前的雇主没有把心理健康看得和身体健康一样重要", "以前的部分雇主把心理健康看得和身体健康一样重要", "我不知道以前的雇主是否把心理健康看得和身体健康一样重要"],
        ["听说过或者观察到有心理健康问题的同事带来的负面影响", "没有听说过或者观察到有心理健康问题的同事带来的负面影响", "可能听说过或者观察到有心理健康问题的同事带来的负面影响"],
        ["愿意在面试中向潜在雇主提出身体健康问题", "不愿意在面试中向潜在雇主提出身体健康问题"],
        ["愿意在面试中向潜在雇主提出心理健康问题", "不愿意在面试中向潜在雇主提出心理健康问题"],
        ["认为会影响职业生涯", "不认为会影响职业生涯", "可能会影响职业生涯", "不会影响职业生涯", "会影响职业生涯"],
        ["会消极的看待", "认为会消极看待", "不会消极的看待", "认为不会消极看待", "可能会消极看待"],
        ["非常愿意分享", "一般般愿意分享", "没有不愿意也没有愿意分享", "一般般不愿意分享", "非常不愿意分享", "我没有精神疾病"],
        ["经历过", "观察过", "没有观察到或经历过", "我不确定是否观察到或经历过"],
        ["有家族病史", "没有家族病史", "我不确定有没有家族病史"],
        ["过去有精神疾病", "过去没有精神疾病", "过去可能有精神疾病"],
        ["目前有精神疾病", "目前没有精神疾病", "目前可能有精神疾病"],
        ["被诊断出有精神健康问题", "没有被诊断出有精神健康问题"],
        ["寻求过心理健康专家的治疗", "没有寻求过心理健康专家的治疗"],
        ["经常会干扰工作", "可能会干扰工作", "很少会干扰工作", "从来没有干扰工作", "不适用"],
        ["经常会影响工作", "可能会影响工作", "很少会影响工作", "从来没有影响工作", "不适用"],
        ["男性", "女性"],
        ["后端开发者", "后端开发者|开发者倡导者/拥护者", "后端开发者|开发者倡导者/拥护者|主管/团队领导", "后端开发者|DevOps/系统管理员", "后端开发者|DevOps/系统管理员|开发者倡导者/拥护者|主管/团队领导", "后端开发者|DevOps/系统管理员|主管/团队领导", "后端开发者|DevOps/系统管理员|主管/团队领导|执行领导", "后端开发者|DevOps/系统管理员|主管/团队领导|其他", "后端开发者|前端开发者", "后端开发者|前端开发者|设计师", "后端开发者|前端开发者|个人开发者", "后端开发者|个人开发者", "后端开发者|主管/团队领导", "后端开发者|主管/团队领导|其他", "后端开发者|支持|DevOps/系统管理员", "后端开发者|支持|主管/团队领导", "设计师", "设计师 | 前端开发人员", "设计师 | 前端开发人员 | 后端开发人员", "设计师 | 前端开发人员 | 后端开发人员 | DevOps/SysAdmin", "设计师 | 前端开发人员 | 后端开发人员 | DevOps/SysAdmin | 其他", "设计师 | 前端开发人员 | 后端开发人员 | 执行领导 | 其他", "设计师 | 前端开发人员 | 后端开发人员 | 其他", "设计师 | 前端开发人员 | 后端开发人员 | 销售 | 督导/团队领导", "设计师 | 前端开发人员 | 后端开发人员 | 督导/团队领导", "设计师 | 前端开发人员 | 后端开发人员 | 督导/团队领导 | 执行领导", "设计师 | 前端开发人员 | 后端开发人员 | 支持 | DevOps/SysAdmin | 督导/团队领导", "设计师 | 个体经营者", "设计师 | 督导/团队领导", "设计师 | 支持 | 督导/团队领导"],
        ["远程上班", "不远程上班"],
    ]

    best_indices = []  # 创建数组存储最佳匹配文本的索引

    for i in range(len(user_contents)):
        text_group = {}
        text_group["text1"] = user_contents[i]
        if i < len(json_data):
            text_group["text2_options"] = json_data[i]
        else:
            text_group["text2_options"] = []
        text_groups.append(text_group)

    for group in text_groups:
        text1 = group["text1"]
        text2_options = group["text2_options"]
        match_scores = []

        max_match_score = -1
        best_index = None

        for index, text2 in enumerate(text2_options):
            inputs = tokenizer(text1, text2, return_tensors='pt', padding=True, truncation=True)
            outputs = model(**inputs)
            logits = outputs.logits
            probabilities = torch.softmax(logits, dim=1)
            match_score = probabilities[:, 1].item()
            match_scores.append((index, text2, match_score))
            if match_score > max_match_score:
                max_match_score = match_score
                best_text2 = text2
                best_index = index

        # print(f"文本1: {text1}")
        # print(f"最佳匹配文本2: {best_text2}")
        if best_index is not None:
            # print(f"选项编号: {best_index + 1}")
            best_indices.append(best_index)  # 将最佳匹配文本的索引添加到数组中
        # for index, text2, match_score in match_scores:
        #     print(f"序号: {index + 1}, 选项: {text2}，匹配得分: {match_score}")
        # print()

    name_data = [
                "self_empl_flag",
                "comp_no_empl",
                "tech_comp_flag",
                "mh_coverage_flag",
                "mh_coverage_awareness_flag",
                "mh_employer_discussion",
                "mh_resources_provided",
                "mh_anonimity_flag",
                "mh_medical_leave",
                "mh_discussion_neg_impact",
                "ph_discussion_neg_impact",
                "mh_discussion_cowork",
                "mh_discussion_supervis",
                "mh_eq_ph_employer",
                "mh_conseq_coworkers",
                "prev_employers_flag",
                "prev_mh_benefits",
                "prev_mh_benefits_awareness",
                "prev_mh_discussion",
                "prev_mh_resources",
                "prev_mh_anonimity",
                "prev_mh_discuss_neg_conseq",
                "prev_ph_discuss_neg_conseq",
                "prev_mh_discussion_cowork",
                "prev_mh_discussion_supervisor",
                "prev_mh_importance_employer",
                "prev_mh_conseq_coworkers",
                "future_ph_specification",
                "why/ why_not",
                "future_mh_specification",
                "why/ why_not2",
                "mh_hurt_on_career",
                "mh_neg_view_cowork",
                "mh_sharing_friends/fam_flag",
                "mh_bad_response_workplace",
                "mh_family_hist",
                "mh_disorder_past",
                "mh_disorder_current",
                "mh_diagnos_proffesional",
                "mh_sought_proffes_treatm",
                "mh_eff_treat_impact_on_work",
                "mh_not_eff_treat_impact_on_work",
                "sex",
                "province_live",
                "city_live",
                "province_work",
                "city_work",
                "work_position",
                "remote_flag",
                ]

    # 跳过的特定值列表
    skip_values = ["why/ why_not", "why/ why_not2", "province_live", "city_live", "province_work", "city_work",]
    # 跳过的特定值数量
    skip_count = len(skip_values)
    skip_counter = 0
    merged_data = {}
    for i, name in enumerate(name_data):
        # 如果值在跳过的特定值列表中，则跳过
        if name in skip_values:
            skip_counter += 1
            if skip_counter == skip_count:
                continue
        if i < len(best_indices):
            merged_data[name] = '0'
        else:
            merged_data[name] = '1'
    # 将字典转换为JSON格式的字符串
    Jdata = json.dumps(merged_data, ensure_ascii=False)
    # print(f"Process中的Jdata:\n{Jdata}\n")
    return Jdata  # 返回存储最佳匹配文本索引的数组

if __name__ == "__main__":
    data = [
        {"role": "user", "content": "这是一段废话"},
        {"role": "user", "content": "我不是个体户, 我在饭店上班"},
        {"role": "user", "content": "我的公司一共88人"},
        {"role": "system", "content": "这是系统回复的内容1"},
        {"role": "user", "content": "我在饭店上班"}
    ]

    best_indices = process_data(data)
    print("最佳匹配文本的索引数组:", best_indices)
